package com.itxiuyixiu.tools.security.ddos;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.itxiuyixiu.tools.utils.back.ReturnMessage;
import com.itxiuyixiu.tools.utils.cache.ScheduleHashMapImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.TimeUnit;

/**
 * 防止爬虫、DDos攻击的过滤器
 *
 * @author 黄磊
 * @date 2020/2/7
 */
public class IpHandlerInterceptor implements HandlerInterceptor {
    @Value("${spring.profiles.active:dev}")
    private String active;

    private Logger log = LoggerFactory.getLogger(getClass());
    private final ObjectMapper objectMapper;
    private static ScheduleHashMapImpl<String, Integer> hashMap = new ScheduleHashMapImpl<>();

    public IpHandlerInterceptor(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    /**
     * 处理请求
     **/
    @Override
    public boolean preHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o) throws Exception {
        if ("dev".equals(active)) {
            return true;
        }
        String ip = httpServletRequest.getRemoteAddr();
        long flag = saveIpAndGetTimes(ip);
        if (flag > 200) {
            refuseVisitReturn(ip, httpServletResponse);
            return false;
        } else {
            return true;
        }
    }

    private void refuseVisitReturn(String ip, HttpServletResponse httpServletResponse) throws Exception {
        httpServletResponse.setStatus(403);
        httpServletResponse.setHeader("Content-type", "text/html;charset=UTF-8");
        httpServletResponse.setCharacterEncoding("UTF-8");
        httpServletResponse.getWriter().write(objectMapper.writeValueAsString(ReturnMessage.getFailureReturn("ip访问次数过多拒绝你的访问")));
        log.error(ip + "因为访问次数过多被限制访问了");
    }

    private long saveIpAndGetTimes(String ip) {
        if (hashMap.get(ip) == null) {
            hashMap.put(ip, 1, 15, TimeUnit.SECONDS);
        } else {
            hashMap.put(ip, hashMap.get(ip) + 1, 15, TimeUnit.SECONDS);
        }
        return hashMap.get(ip);
    }
}
